# coding: utf-8
# implementation of GA for TSP
import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import sample, randint, shuffle, uniform
from time import time, gmtime, strftime, localtime
import argparse
from icecream import ic
from tqdm import trange

from TSP import TSP

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)

parser = argparse.ArgumentParser()
parser.add_argument("--case", required=True, type=str, help="The case name!")
args = parser.parse_args()

case = args.case

instance = TSP(path=f'../../data/TSPLIB/{case}.tsp.txt')
Distance = instance.distance_matrix  # distance matrix
V = deepcopy(instance.cities)  # city set


def make_dir(path):
    path = path.strip().rstrip("\\")
    if os.path.exists(path):
        print(path + ' already exists.')
        return False
    else:
        os.makedirs(path)
        return True


def find_nearest(city, candidates):
    # find the nearest neighbor of 'city' among 'candidates'
    tar, min_dis = -1, 1e99
    for c in candidates:
        if min_dis > Distance[city][c]:
            min_dis = Distance[city][c]
            tar = c
    assert tar != -1
    return tar


class Solution:
    def __init__(self, init='RND'):
        self.route = None  # route array
        self.length = None  # route length
        if init == 'EMP':
            # empty route
            self.route = []
            return
        elif init == 'RND':
            # random initialize route
            self.route = deepcopy(V)
            shuffle(self.route)
        elif init == 'GRD':
            # greedily initialize route
            depot = sample(V, k=1)[0]
            cities = deepcopy(V)
            cities.remove(depot)  # cities exclude depot
            self.route = [depot]
            while cities:
                self.route.append(find_nearest(city=self.route[-1], candidates=cities))
                cities.remove(self.route[-1])
        else:
            print(f'error: init={init}. ')
            quit(-1)
        self.calculate_length()
        self.validate()

    def calculate_length(self):
        self.length = 0
        for i in range(1, len(self.route)):
            self.length += Distance[self.route[i - 1]][self.route[i]]
        self.length += Distance[self.route[-1]][self.route[0]]  # add the depot to form the circle

    def validate(self):
        diff1 = set(V) - set(self.route)
        diff2 = set(self.route) - set(V)
        assert len(diff1) == 0 & len(diff2) == 0


def _2_OPT(route):
    improve = 0
    for i in range(1, len(route)):
        for j in range(i + 1, len(route)):
            d1 = Distance[route[i - 1]][route[i]] + Distance[route[j - 1]][route[j]]
            d2 = Distance[route[i - 1]][route[j - 1]] + Distance[route[i]][route[j]]
            if d2 < d1:
                route = route[:i] + list(reversed(route[i:j])) + route[j:]
                improve += d1 - d2
    return route, improve


def rotate(solution):
    # random rotate the segments between 'i' and 'j' in the route of solution
    indexes = list(range(0, len(solution.route)))
    i, j = sorted(sample(indexes, 2))
    new_solution = deepcopy(solution)
    new_solution.route = solution.route[0:i] + solution.route[j:i - 1:-1] + solution.route[j:]
    # print(new_solution.route)
    new_solution.calculate_length()
    new_solution.validate()
    return new_solution if new_solution.length < solution.length else solution


def partially_mapped_crossover(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = list(range(0, len(P1.route)))
    a, b = sorted(sample(indexes, 2))
    for i in range(a, b + 1):
        O1.route[i], O2.route[i] = O2.route[i], O1.route[i]  # swap segments
    for i in range(len(O1.route)):
        if i < a or i > b:
            while O1.route[i] in O1.route[a:b + 1]:
                repeat_index = O1.route[a:b + 1].index(O1.route[i])
                O1.route[i] = O2.route[a:b + 1][repeat_index]  # replace i with allele in O2
            while O2.route[i] in O2.route[a:b + 1]:
                repeat_index = O2.route[a:b + 1].index(O2.route[i])
                O2.route[i] = O1.route[a:b + 1][repeat_index]
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def order_crossover(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = list(range(0, len(P1.route)))
    a, b = sorted(sample(indexes, 2))
    seq1, seq2 = [], []
    for i in list(range(b + 1, len(P1.route))) + list(range(0, b + 1)):
        if P2.route[i] not in O1.route[a:b + 1]:
            seq2.append(P2.route[i])
        if P1.route[i] not in O2.route[a:b + 1]:
            seq1.append(P1.route[i])
    j = 0
    for i in list(range(b + 1, len(P1.route))) + list(range(0, a)):
        O1.route[i] = seq2[j]
        O2.route[i] = seq1[j]
        j = j + 1
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def cycle_crossover(P1, P2):
    O1, O2 = deepcopy(P1), deepcopy(P2)
    indexes = [_ for _ in range(len(P1.route))]
    first = sample(list(range(len(P1.route))), k=1)[0]  # choose the first city in P1.route to inherit
    # first = 0
    i = first
    while P2.route[i] != O1.route[first]:
        indexes[i] = -1  # record the inherited index of P1.route
        i = P1.route.index(P2.route[i])  # index of the allele in P1.route
    indexes[i] = -1
    for i in indexes:
        if i != -1:
            O1.route[i] = P2.route[i]
    indexes = [_ for _ in range(len(P2.route))]
    first = sample(list(range(len(P2.route))), k=1)[0]  # choose the first city in P1.route to inherit
    i = first
    while P1.route[i] != O2.route[first]:
        indexes[i] = -1  # record the inherited index of P1.route
        i = P2.route.index(P1.route[i])  # index of the allele in P1.route
    indexes[i] = -1
    for i in indexes:
        if i != -1:
            O2.route[i] = P1.route[i]
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def cycle_crossover_2(P1, P2):
    # Step1. Choose two parents for mating.
    O1, O2 = deepcopy(P1), deepcopy(P2)
    # Step2. Select 1st bit from P2 as a 1st bit of O1.
    O1.route[0] = P2.route[0]
    # Step3. Select


def edge_recombination_crossover(P1, P2):
    """
        Most operators take into account the position and the order of the city.
        This operator takes into account the links between these cities.
    """
    O1, O2 = deepcopy(P1), deepcopy(P2)
    n = len(P1.route)
    edge_list = [set() for _ in range(n + 1)]
    for i in range(n):
        edge_list[P1.route[i]].add(P1.route[(i - 1) % n])
        edge_list[P1.route[i]].add(P1.route[(i + 1) % n])
        edge_list[P2.route[i]].add(P2.route[(i - 1) % n])
        edge_list[P2.route[i]].add(P2.route[(i + 1) % n])
    # for i in range(1, len(edge_list)):
    #     print(f'edge {i}: ', edge_list[i])
    # build O1
    cities = deepcopy(V)
    ini = sample(cities, k=1)[0]
    O1.route = [ini]
    cities.remove(ini)
    while cities:
        li = [i for i in edge_list[O1.route[-1]] if i not in O1.route]
        if li:
            c = min(li, key=lambda x: len(edge_list[x]))
        else:
            c = sample(cities, k=1)[0]
        O1.route.append(c)
        cities.remove(c)
    # build O2
    cities = deepcopy(V)
    ini = sample(cities, k=1)[0]
    O2.route = [ini]
    cities.remove(ini)
    while cities:
        li = [i for i in edge_list[O2.route[-1]] if i not in O2.route]
        shuffle(li)
        if li:
            c = min(li, key=lambda x: len(edge_list[x]))
        else:
            c = sample(cities, k=1)[0]
        O2.route.append(c)
        cities.remove(c)
    O1.calculate_length()
    O2.calculate_length()
    O1.validate()
    O2.validate()
    return O1, O2


def alternating_edge_crossover(P1, P2):
    pass


def sequential_constructive_crossover(P1, P2):
    pass


def mutate(solution):
    indexes = list(range(0, len(solution.route)))
    i, j = sorted(sample(indexes, 2))
    solution.route[i], solution.route[j] = solution.route[j], solution.route[i]
    solution.calculate_length()
    return solution


plt.figure(figsize=(10, 4))
max_epoch = 1
max_gen = 10000
max_eva = 1000 * instance.n
NP = 50  # size of population
Pc = 0.7  # cross prob
Pm = 0.1  # mutation prob
# operators
PMX = partially_mapped_crossover
OX = order_crossover
CX = cycle_crossover
CX2 = cycle_crossover_2
ER = edge_recombination_crossover


if __name__ == '__main__':
    date = strftime("%Y-%m-%d", localtime())
    save_path = f'./results/{date}/{case}'
    make_dir(save_path)
    start_time = time()
    fitness = []
    for epoch in trange(1, max_epoch + 1):
        POP = []
        # Initialization.
        for _ in range(NP):
            POP.append(Solution(init='RND'))
            # for i in range(1000):
            #     plt.clf()
            #     instance.draw(plt, route=POP[-1].route + [POP[-1].route[0]], show_text=True)
            #     plt.title(f'{i}: {POP[-1].length:.2f}')
            #     plt.pause(0.001)
            #     POP[-1].route, improve = _2_OPT(POP[-1].route)
            #     # print(POP[-1].route, POP[-1].length, improve)
            #     POP[-1].length -= improve

        # Progress
        pBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest_list = []
        pBest_list = []
        pWorst_list = []
        for gen in trange(1, max_gen + 1):
            # Select.
            offspring = []
            fit = np.array([1 / (1 + POP[i].length) for i in range(len(POP))])  # fitness = 1 / (1 + length)
            fit = fit / fit.sum()  # fitness proportion
            winners = np.random.choice(a=np.arange(0, len(POP)), size=NP, replace=True,
                                       p=fit)  # roulette wheel selection
            for winner in winners:  # add winners to offspring population
                offspring.append(POP[winner])
            offspring[offspring.index(max(offspring, key=lambda _: _.length))] = deepcopy(gBest)  # elite strategy
            POP = offspring  # replace the old POP
            # Crossover.
            shuffle(POP)
            for i in range(1, NP, 2):
                if uniform(0, 1) < Pc:
                    POP[i - 1], POP[i] = ER(POP[i - 1], POP[i])
            # Mutation.
            for i in range(1, NP):
                if uniform(0, 1) < Pm:
                    POP[i] = mutate(POP[i])
            # Record gBest
            pBest = deepcopy(min(POP, key=lambda _: _.length))
            if gBest.length > pBest.length:
                gBest = deepcopy(pBest)
            pBest_list.append(pBest.length)
            gBest_list.append(gBest.length)
            pWorst_list.append(max(POP, key=lambda _: _.length).length)
            print(f'method=GA, case={case}, NP={NP}, Pc={Pc}, Pm={Pm}. '
                  f'At epoch {epoch} gen {gen}: {round(gBest.length, 2)}', end='\r')
            plt.clf()
            plt.subplot(1, 2, 1)
            plt.plot(range(len(gBest_list)), gBest_list, label='gBest')
            # plt.plot(range(len(pBest_list)), pBest_list, label='pBest')
            # plt.plot(range(len(pWorst_list)), pWorst_list, label='pWorst')
            plt.legend()
            plt.title(f'total length: {gBest.length:.2f}')
            plt.subplot(1, 2, 2)
            instance.draw(plt, route=gBest.route + [gBest.route[0]])
            # plt.pause(0.001)
        plt.show()

        fitness.append(gBest_list[-1])
        mode = 'w' if epoch == 1 else 'ab'
        # ---------> Change save file here! <---------
        with open(save_path + rf'/TSP_GA_CX_rand.txt', mode) as f:
            np.savetxt(f, np.array(gBest_list).reshape(1, len(gBest_list)))

    print(f'method=GA, case={case}, NP={NP}, Pc={Pc}, Pm={Pm}. '
          f'At epoch {max_epoch} gen {max_gen}: {round(np.array(fitness).mean(), 2)}')
    t = strftime('%H:%M:%S', gmtime(time() - start_time))
    print(f'Time cost {t}')
